package first.group.ssh;

import first.group.ssh.authorize.Authorizer;
import first.group.ssh.channel.Channel;
import first.group.ssh.channel.ChannelMultiplexer;
import first.group.ssh.encrypt.Encryption;
import first.group.ssh.globalstate.GlobalStateManager;
import first.group.ssh.keyexchange.KeyAlgoNegotiator;
import first.group.ssh.keyexchange.KeyExchangeMethod;
import first.group.ssh.keyexchange.KeyExchanger;
import first.group.ssh.keyexchange.KeyGenerator;
import first.group.ssh.mac.MacAlgorithm;
import first.group.ssh.packet.DebugPacket;
import first.group.ssh.packet.DisconnectPacket;
import first.group.ssh.packet.GlobalRequestFailurePacket;
import first.group.ssh.packet.GlobalRequestPacket;
import first.group.ssh.packet.GlobalRequestSuccessPacket;
import first.group.ssh.packet.IgnorePacket;
import first.group.ssh.packet.Packet;
import first.group.ssh.packet.PacketFactory;
import first.group.ssh.packet.PacketSerializer;
import first.group.ssh.packet.UnimplementedPacket;
import first.group.ssh.readBuffer.ReadBuffer;
import first.group.ssh.version.Version;
import first.group.ssh.version.VersionExchanger;
import first.group.ssh.writebuffer.ByteBufferContent;
import first.group.ssh.writebuffer.WriteBuffer;
import src.AESCipherCTR;

import java.io.EOFException;
import java.io.IOException;
import java.math.BigInteger;
import java.net.InetSocketAddress;
import java.nio.channels.SelectionKey;
import java.nio.channels.SocketChannel;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.util.LinkedList;
import java.util.Random;
import java.util.logging.Level;
import java.util.logging.Logger;
import javax.crypto.Cipher;
import javax.crypto.Mac;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;

public class Connection implements ConnectionCallback, KeyExchangeMethod.Callback {

    private static final Logger logger = Logger.getLogger("Connection");

    private Random random;

    private final PacketFactory packetFactory = new PacketFactory();
    private final PacketSerializer packetSerializer = new PacketSerializer();

    private final Version clientVersion = Version.VERSION_STRING;
    private final VersionExchanger versionExchanger = new VersionExchanger(this, clientVersion);
    private final Authorizer authorizer = new Authorizer(this);
    private final KeyAlgoNegotiator keyAlgoNegotiator;
    private KeyExchanger keyExchanger;
    private final ChannelMultiplexer channelMultiplexer = new ChannelMultiplexer(this);
    private final GlobalStateManager globalStateManager = new GlobalStateManager(this);

    private BigInteger sharedSecret;
    private byte[] exchangeHash;
    private byte[] sessionId;

    private final WriteBuffer writeBuffer = new WriteBuffer();
    private ReadBuffer readBuffer = new ReadBuffer();

    public final InetSocketAddress address;

    private void initRandom() {
        try {
            random = SecureRandom.getInstanceStrong();
        } catch (NoSuchAlgorithmException ex) {
            logger.warning("Unable to use strong secure random, fallback to regular random");
            random = new Random();
        }
    }

    public Connection(InetSocketAddress address) {
        this.address = address;
        initRandom();
        keyAlgoNegotiator = new KeyAlgoNegotiator(random, this);
    }

    public void setAuthorizeListener(Authorizer.AuthorizeListener listener) {
        authorizer.setAuthorizeListener(listener);
    }

    @Override
    public void finish() {
        switch (status) {
            case DISCONNECTED, DISCONNECTING, FAILED -> {
                throw new IllegalStateException("Finish in state " + status.name());
            }
            case CONNECTING -> {
                setStatus(Status.VERSION_EXCHANGING);
                versionExchanger.init();
            }
            case VERSION_EXCHANGING -> {
                setStatus(Status.KEY_ALGO_NEGOTIATING);
                packetFactory.fillIn(readBuffer);
                readBuffer = null;
                keyAlgoNegotiator.init();
            }
            case KEY_ALGO_NEGOTIATING -> {
                setStatus(Status.KEY_EXCHANGING);
                keyExchanger = new KeyExchanger(
                        clientVersion,
                        versionExchanger.getServerVersion(),
                        keyAlgoNegotiator.getClientInitPacket(),
                        keyAlgoNegotiator.getServerInitPacket(),
                        random,
                        keyAlgoNegotiator.getSuite(),
                        this,
                        this
                );
                keyExchanger.init();
            }
            case KEY_EXCHANGING -> {
                applyNewKey();
                setStatus(Status.AUTHORIZING);
                authorizer.init();
            }
            case AUTHORIZING -> {
                setStatus(Status.AUTHORIZED);
            }
            case AUTHORIZED -> {
                throw new IllegalStateException("Channel should call finishChannel(Channel) instead");
            }
            default ->
                throw new AssertionError(status.name());
        }
    }

    @Override
    public WriteBuffer getWriteBuffer() {
        return writeBuffer;
    }

    @Override
    public void appendPacket(Packet packet) {
        var data = packetSerializer.serialize(packet);
        writeBuffer.append(new ByteBufferContent(data));
        wakeup();
    }

    public enum Status {
        DISCONNECTED,
        CONNECTING,
        VERSION_EXCHANGING,
        KEY_ALGO_NEGOTIATING,
        KEY_EXCHANGING,
        AUTHORIZING,
        AUTHORIZED,
        DISCONNECTING,
        FAILED
    };

    private Status status = Status.DISCONNECTED;
    private final LinkedList<StatusListener> statusListeners = new LinkedList<>();

    public interface StatusListener {

        void statusChanged(Status status);
    }

    public Status getStatus() {
        return status;
    }

    public void addStatusListener(StatusListener listener) {
        statusListeners.add(listener);
    }

    public void removeStatusListener(StatusListener listener) {
        statusListeners.remove(listener);
    }

    public void setStatus(Status status) {
        if (this.status != status) {
            logger.log(Level.FINE, "Connection {0} status changed to {1}", new Object[]{this, status.name()});
            this.status = status;
            for (var listener : statusListeners) {
                listener.statusChanged(status);
            }
        }
    }

    public void disconnect() {
        try {
            // TODO wait for write disconnect packet
            socketChannel.close();
        } catch (IOException ex) {
            throw new RuntimeException(ex);
        }
        setStatus(Status.DISCONNECTING);
    }

    private void doFinal(SelectionKey key) {
        logger.info("Connection disconnected");
        if (status == Status.DISCONNECTING) {
            setStatus(Status.DISCONNECTED);
        } else {
            key.cancel();
        }
        ConnectionService.getInstance().removeConnection(this);
    }

    private Exception failedException = null;

    private void fail(Exception ex, boolean logException) {
        if (logException) {
            logger.log(Level.WARNING, "Connection failed", ex);
        } else {
            logger.log(Level.WARNING, "Connection failed");
        }
        try {
            socketChannel.close();
        } catch (IOException closeException) {
            logger.log(Level.SEVERE, "close failed", closeException);
        }
        failedException = ex;
        setStatus(Status.FAILED);
    }

    public Exception getFailedException() {
        return failedException;
    }

    private SocketChannel socketChannel = null;

    public SocketChannel openSocketChannel() {
        if (socketChannel != null) {
            throw new IllegalStateException("Already opened a socket channel");
        }
        try {
            socketChannel = SocketChannel.open();
            socketChannel.configureBlocking(false);
            socketChannel.connect(address);
            setStatus(Status.CONNECTING);
        } catch (IOException ex) {
            logger.log(Level.SEVERE, "Exception when connecting: ", ex);
            fail(ex, false);
            return null;
        }
        return socketChannel;
    }

    @Override
    public void wakeup() {
        ConnectionService.getInstance().wakeup();
    }

    private void checkConnectable(SelectionKey key) {
        if (status == Status.CONNECTING) {
            key.interestOpsOr(SelectionKey.OP_CONNECT);
        } else {
            key.interestOpsAnd(~SelectionKey.OP_CONNECT);
        }
    }

    private boolean shouldRead() {
        switch (status) {
            case DISCONNECTED -> {
                return false;
            }
            case CONNECTING -> {
                return false;
            }
            case VERSION_EXCHANGING -> {
                return readBuffer.hasRemaining() && versionExchanger.shouldRead();
            }
            case KEY_ALGO_NEGOTIATING -> {
                return keyAlgoNegotiator.shouldRead();
            }
            case KEY_EXCHANGING -> {
                return keyExchanger.shouldRead();
            }
            case AUTHORIZING -> {
                return authorizer.shouldRead();
            }
            case AUTHORIZED -> {
                return true;
            }
            case DISCONNECTING -> {
                return false;
            }
            case FAILED -> {
                logger.warning("Checking reading state after failing");
                return false;
            }
        }
        throw new AssertionError();
    }

    public void checkState(SelectionKey key) {
        if (!key.isValid()) {
            doFinal(key);
            return;
        }
        checkConnectable(key);
        if (shouldRead()) {
            key.interestOpsOr(SelectionKey.OP_READ);
        } else {
            key.interestOpsAnd(~SelectionKey.OP_READ);
        }
        if (writeBuffer.shouldWrite()) {
            key.interestOpsOr(SelectionKey.OP_WRITE);
        } else {
            key.interestOpsAnd(~SelectionKey.OP_WRITE);
        }
    }

    private void processReadData() throws IOException {
        if (status == Status.VERSION_EXCHANGING) {
            while (readBuffer.hasData() && shouldRead()) {
                versionExchanger.read(readBuffer);
                if (status != Status.VERSION_EXCHANGING) {
                    break;
                }
            }
        } else {
            while (true) {
                if (!shouldRead()) {
                    break;
                }
                Packet packet = packetFactory.getPacket();
                if (packet == null) {
                    break;
                }
                if (packet instanceof IgnorePacket) {
                    continue;
                }
                if (packet instanceof DebugPacket debugPacket) {
                    if (debugPacket.alwaysDisplay) {
                        logger.log(Level.INFO, "Debug packet: {0}", debugPacket.getMessage());
                    } else {
                        logger.log(Level.FINE, "Debug packet: {0}", debugPacket.getMessage());
                    }
                    continue;
                }
                if (packet instanceof DisconnectPacket) {
                    logger.log(Level.INFO, "Service requested disconnect");
                    disconnect();
                    return;
                }
                if (packet instanceof UnimplementedPacket unimplementedPacket) {
                    // TODO how to proper handle this packet
                    logger.log(Level.WARNING, "Packet seq {0} received UnimplementedPacket", unimplementedPacket.seq);
                    return;
                }
                switch (status) {
                    case KEY_ALGO_NEGOTIATING -> {
                        keyAlgoNegotiator.read(packet);
                    }
                    case KEY_EXCHANGING -> {
                        keyExchanger.read(packet);
                    }
                    case AUTHORIZING -> {
                        authorizer.read(packet);
                    }
                    case AUTHORIZED -> {
                        if (packet instanceof GlobalRequestPacket
                                || packet instanceof GlobalRequestSuccessPacket
                                || packet instanceof GlobalRequestFailurePacket) {
                            globalStateManager.read(packet);
                        } else {
                            channelMultiplexer.read(packet);
                        }
                    }
                    case DISCONNECTED, DISCONNECTING, FAILED -> {
                        logger.log(Level.WARNING, "Read data when {0}", status.name());
                    }
                    case CONNECTING, VERSION_EXCHANGING ->
                        throw new AssertionError(status.name());
                    default ->
                        throw new AssertionError(status.name());
                }
            }
        }
    }

    private boolean processRead(SelectionKey key) throws IOException {
        if (status == Status.VERSION_EXCHANGING) {
            while (key.isReadable()) {
                if (readBuffer.hasRemaining()) {
                    var len = readBuffer.fillData(socketChannel);
                    if (len <= 0) {
                        break;
                    }
                    processReadData();
                    if (status != Status.VERSION_EXCHANGING) {
                        break;
                    }
                } else {
                    processReadData();
                    if (status != Status.VERSION_EXCHANGING) {
                        break;
                    }
                    return false;
                }
            }
        } else {
            while (key.isReadable()) {
                var len = packetFactory.fillData(socketChannel);
                if (len <= 0) {
                    break;
                }
            }
            processReadData();
        }
        return true;
    }

    private boolean ensureAlive(SelectionKey key) {
        if (!key.isValid()) {
            doFinal(key);
            return false;
        } else if (status == Status.DISCONNECTING || status == Status.DISCONNECTED || status == Status.FAILED) {
            return false;
        }
        return true;
    }

    public void processSelection(SelectionKey key) {
        if (!ensureAlive(key)) {
            return;
        }
        if (key.isConnectable()) {
            try {
                socketChannel.finishConnect();
                finish();
            } catch (IOException ex) {
                key.cancel();
                fail(ex, true);
                return;
            }
        }
        if (!ensureAlive(key)) {
            return;
        }
        try {
            processRead(key);
        } catch (EOFException ex) {
            doFinal(key);
            return;
        } catch (IOException | RuntimeException ex) {
            key.cancel();
            logger.log(Level.SEVERE, "Exception in reading", ex);
            fail(ex, false);
            return;
        }
        if (!ensureAlive(key)) {
            return;
        }
        if (key.isWritable()) {
            try {
                writeBuffer.doWrite(socketChannel);
            } catch (IOException ex) {
                key.cancel();
                logger.log(Level.SEVERE, "Exception in writing", ex);
                fail(ex, false);
            }
        }
    }

    private Cipher initCipher(Encryption encryption, byte[] key, byte[] iv) throws InvalidKeyException, InvalidAlgorithmParameterException {
        var cipher = encryption.newCipher();
        var keySpec = new SecretKeySpec(key, encryption.getKeyAlgorithm());
        var ivSpec = new IvParameterSpec(iv);
        cipher.init(Cipher.ENCRYPT_MODE, keySpec, ivSpec);
        return cipher;
    }

    private Mac initMac(MacAlgorithm algorithm, byte[] key) throws InvalidKeyException, InvalidKeyException {
        var mac = algorithm.newMac();
        var keySpec = new SecretKeySpec(key, mac.getAlgorithm());
        mac.init(keySpec);
        return mac;
    }

    private void applyNewKey() {
        try {
            var suite = keyAlgoNegotiator.getSuite();
            var method = suite.keyExchangeMethod();
            var pack = KeyGenerator.pack(sharedSecret, exchangeHash);

            var c2sMacAlgo = suite.clientToServerMac();
            var s2cMacAlgo = suite.serverToClientMac();

            var c2sIv = KeyGenerator.generate(method, pack, (byte) 'A', sessionId, 16);
            var s2cIv = KeyGenerator.generate(method, pack, (byte) 'B', sessionId, 16);
            var c2sKey = KeyGenerator.generate(method, pack, (byte) 'C', sessionId, 16);
            var s2cKey = KeyGenerator.generate(method, pack, (byte) 'D', sessionId, 16);
            var c2sMacKey = KeyGenerator.generate(method, pack, (byte) 'E', sessionId, c2sMacAlgo.keyLength());
            var s2cMacKey = KeyGenerator.generate(method, pack, (byte) 'F', sessionId, s2cMacAlgo.keyLength());

            var c2sCipher = new AESCipherCTR(c2sKey, c2sIv);
            var s2cCipher = new AESCipherCTR(s2cKey, s2cIv);
            var c2sMac = initMac(c2sMacAlgo, c2sMacKey);
            var s2cMac = initMac(s2cMacAlgo, s2cMacKey);

            packetFactory.setCipher(s2cCipher, s2cMac);
            packetSerializer.setCipher(c2sCipher, c2sMac);
        } catch (InvalidKeyException ex) {
            throw new RuntimeException(ex);
        }
    }

    @Override
    public void onKeyExchanged(BigInteger secret, byte[] sessionId) {
        this.sharedSecret = secret;
        if (this.exchangeHash == null) {
            this.exchangeHash = sessionId;
        }
        this.sessionId = sessionId;
        keyExchanger.prepareFinish();
    }

    @Override
    public void onKeyExchangeFailed() {
        if (status != Status.KEY_EXCHANGING) {
            throw new IllegalStateException("Key exchange must happen in state onKeyExchangeFailed, but now is " + status.name());
        }
        logger.warning("Key exchange failed");
        // TODO send a failed packet before real failing
        fail(null, false);
    }

    public void openChannel(Channel channel) {
        channelMultiplexer.openChannel(channel);
    }
}
